'''
Created on Jul 25, 2010

@author: oabalbin
'''
import sys
import numpy as np
import MySQLdb;
from collections import deque, defaultdict

from interactome.common.classes import interaction


class knownGene():
    def __init__(self, kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description):
        self.kgID, self.mRNA, self.spID, self.spDisplayID, self.geneSymbol, self.refseq, self.protAcc, self.description=\
        kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description
        self.prot_length = 0
    
    def update_protein_length(self,length):
        self.prot_length = length
    

class ucsc_query():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        """ Database Information """
        self.myhost = dbhost
        self.myuser = dbuser
        self.mypasswd = dbpasswd
        self.mydb = database
        self.create_connector()
    
    def create_connector(self):
        self.conn=MySQLdb.connect(host=self.myhost,user=self.myuser,passwd=self.mypasswd,db=self.mydb)
    
    def close_connector(self, connector):
        connector.close
    
    def get_genelist_from_ucsc(self,genelist):
        
        knowngeneslist = defaultdict(deque)
        cursor=self.conn.cursor()
        notfound = set()
        for gen, prots in genelist.iteritems():
            thisprots = prots[0]
            for prot in thisprots:
                query = 'select distinct * from kgXref  where geneSymbol=\"' + gen +'\" and protAcc=\"' + prot +'\";'
                cursor.execute(query)
                 
                # Fetch all the rows in a list of lists.                    
                results = cursor.fetchall()
                if gen == "EPN2":
                    print gen, prot
                    print results
                    
                if not results:
                    notfound.add(gen)
                    continue

                for row in results:
                    kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description =\
                    row[0],row[1],row[2],row[3],row[4],row[5],row[6],row[7]
                                                
                    if mRNA != refseq:
                        continue
                    
                    knowngeneslist[gen].append(knownGene(kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description))
        
        
        genes_not_found = list(notfound.difference(set(knowngeneslist.keys())))
        if  "EPN2" in notfound:
            print "carajo"
        knowngeneslist, notfound = self.get_genelist_from_ucsc_recue(genes_not_found,knowngeneslist)
        

        thediff = list(set(genelist).difference(set(knowngeneslist.keys())))
        #print ",".join(notfound).replace(',','\n')
        print len(genelist)
        print len(knowngeneslist)
        print len(thediff)

        
        return knowngeneslist, notfound
    
    def get_genelist_from_ucsc_recue(self, genes_not_found, knowngeneslist):
        
        cursor=self.conn.cursor()
        notfound = set()    
        
        
        for gen in genes_not_found:
            query = 'select distinct * from kgXref  where geneSymbol=\"' + gen +'\" ;'
            cursor.execute(query)
            results = cursor.fetchall()
                                            
            if not results:
                #knowngeneslist[gen].append(knownGene("NF", "NF", "NF", "NF", gen, "NF", "NF", "NF"))
                notfound.add(gen)
                continue

            for row in results:
                kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description =\
                row[0],row[1],row[2],row[3],row[4],row[5],row[6],row[7]
                
                if gen == "EPN2":
                    print gen, kgID, mRNA, refseq
                    print gen
                    if mRNA != refseq:
                        print "second different"


                knowngeneslist[gen].append(knownGene(kgID, mRNA, spID, spDisplayID, geneSymbol, refseq, protAcc, description))
                
        return knowngeneslist, notfound
    
    def get_prot_length(self,knowngeneslist):
        
        cursor=self.conn.cursor()
        prot_length = {}
        
        for gene, prots in knowngeneslist.iteritems():
            avg=0.0
            for prot in prots:
                query =  'select * from knownGenePep where name =\"' + prot.kgID +'\" ;'
                cursor.execute(query)
                # Fetch all the rows in a list of lists.
                results = cursor.fetchall()
                if not results:
                    print gene, prot.kgID, 
                    continue
                
                seq = results[0][1]
                prot.update_protein_length(len(seq))
                avg += len(seq)
                if gene == "EPN2":
                    print avg, prot.kgID, prot.refseq

            prot_length[gene] = int(round(float(avg)/len(prots)))
            
        return prot_length
    
    
    def print_list_genes(self, knowngeneslist, outfile):
        
        for g, genelist in knowngeneslist.iteritems():
            for gene in genelist:
                if gene.spDisplayID.find("_HUMAN") != -1:
                    line =  [gene.geneSymbol, gene.refseq, str(gene.prot_length),\
                              gene.spDisplayID, gene.protAcc, gene.description]
                    outfile.write(",".join(line).replace(',','\t')+'\n')
    
    def print_list_prots(self, knowngeneslist, notfound, outfile):
        
        for gene, plength in knowngeneslist.iteritems():
            line =[gene, str(plength)]
            if gene == "EPN2":
                print line
            outfile.write(",".join(line).replace(',','\t')+'\n')
        
        for gene in notfound:
            line =[gene, str(1)]
                        
            if gene == "EPN2":
                print line
            outfile.write(",".join(line).replace(',','\t')+'\n')

                
    def list_of_names(self,inputfile):
        
        genlist=defaultdict(deque)
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            genlist[fields[0]].append(fields[1:])
        
        return genlist


"""
This script look for the protein length of each gene in the genelist
"""

myuscs_db = ucsc_query("localhost", "oabalbin", "oscar", "ucsc_tables")
inputfile = '/data/projects/prelims/phosphorylated_proteins.txt'
outfile = '/data/projects/prelims/phosphorylated_proteins_length.txt'

genelist = myuscs_db.list_of_names(open(inputfile))
ucsc_knowgenes, notfound = myuscs_db.get_genelist_from_ucsc(genelist)
prot_length = myuscs_db.get_prot_length(ucsc_knowgenes)
myuscs_db.print_list_prots(prot_length, notfound, open(outfile,'w'))


            
        
        
        
        

        
        

